#include "TriangleMesh.h"

using namespace OBR;

bool __loadOBJFileAsSingleMesh(std::string path, TriangleMesh &mesh)
{
    // check path
    if (path.length() == 0)
    {
        std::cerr << "path provided to OBJLoader cannot be empty" + path << std::endl;
    }

    bool success;
    objl::Loader loader;
    // load file
    success = loader.LoadFile(path);
    if (success)
    {
        int mesh_id = 0;
        for (objl::Mesh &m : loader.LoadedMeshes)
        {
            for (unsigned int &idx : m.Indices)
            {
                objl::Vertex v = m.Vertices[idx];
                Vertex vert;
                vert.n = make_float3(v.Normal.X, v.Normal.Y, v.Normal.Z);
                vert.pos = make_float3(v.Position.X, v.Position.Y, v.Position.Z);
                vert.tex_uv = make_float2(v.TextureCoordinate.X, v.TextureCoordinate.Y);
                mesh.vertexes.push_back(vert);
            }
            for (int base_offset = 0; base_offset < m.Indices.size(); base_offset += 3)
            {
                Triangle tri;
                tri.vertexes[0] = m.Indices[base_offset];
                tri.vertexes[1] = m.Indices[base_offset + 1];
                tri.vertexes[2] = m.Indices[base_offset + 2];
                mesh.triangles.push_back(tri);
            }
            mesh_id++;
        }

        mesh.updateAABB();
        // NOTE: materials for obj files are discareded
    }
    else
    {
        std::cerr << "error loading obj file @ " + path << std::endl;
    }
    return success;
}

std::shared_ptr<CudaProgram> TriangleMesh::program;
CudaBuffer TriangleMesh::sbt_data_payload;

TriangleMesh::TriangleMesh() : Mesh("shaders/Mesh/TriangleMesh.cu")
{
    // init program opt & desc
    this->program_desc.hitgroup.entryFunctionNameCH = "__closesthit__triangle_mesh";
    this->program_desc.kind = OPTIX_PROGRAM_GROUP_KIND_HITGROUP;
    this->program_desc.flags = OPTIX_GEOMETRY_FLAG_NONE;
}

TriangleMesh::~TriangleMesh()
{
}

void TriangleMesh::_buildFromConfig(const Json::Value &config)
{
    if (config["source type"].asString() == "raw")
    {
        // TODO: move this to another function
        int n_vertex = config["num_vertex"].asInt();
        int n_triangle = config["num_triangle"].asInt();
        // vertex
        for (int v = 0; v < n_vertex; v++)
        {
            Vertex vert;
            // position
            vert.pos = ConfigHelpers::loadFloat3(config["vertexes"][v]["position"]);
            // normal (optional)
            if (config["vertexes"][v].isMember("normal"))
            {
                vert.n = ConfigHelpers::loadFloat3(config["vertexes"][v]["n"]);
            }
            else
            {
                vert.n = make_float3(0);
            }
            // uv (optional)
            if (config["vertexes"][v].isMember("uv"))
            {
                vert.tex_uv = ConfigHelpers::loadFloat2(config["vertexes"][v]["uv"]);
            }
            else
            {
                vert.tex_uv = make_float2(0);
            }
            // add vertex
            this->vertexes.push_back(vert);

            // compute AABB
            OptixAabb p = Bounds::fromPoint(vert.pos);
            this->aabb = Bounds::combine(this->aabb, p);
        }
        // triangle
        for (int t = 0; t < n_triangle; t++)
        {
            Triangle tri;
            tri.mesh_id = 0; // TODO: add meshid
            tri.vertexes[0] = config["triangles"][t][0].asInt();
            tri.vertexes[1] = config["triangles"][t][1].asInt();
            tri.vertexes[2] = config["triangles"][t][2].asInt();
            triangles.push_back(tri);
        }
    }
    else if (config["source type"].asString() == "obj file")
    {
        std::string source_path = config["source path"].asString();
        if (source_path.rfind("./", 0) == 0)
        {
            source_path = std::string(PROJ_ROOT) + source_path.substr(1, source_path.size() - 1);
        }
        __loadOBJFileAsSingleMesh(source_path, *this);
        for (Triangle &tri : this->triangles)
        {
            tri.mesh_id = 0; // TODO: add meshid
        }
    }
}

OptixBuildInput TriangleMesh::getASBuildInput()
{
    OptixBuildInput build_input = {};

    // upload to device
    uint vertex_size = sizeof(Vertex) * this->vertexes.size();
    this->d_vertexes.resize(vertex_size);
    this->d_vertexes.upload(this->vertexes.data(), vertex_size);
    uint triangle_size = sizeof(Triangle) * this->triangles.size();
    this->d_triangles.resize(triangle_size);
    this->d_triangles.upload(this->triangles.data(), triangle_size);

    // build flags
    if (build_flags)
    {
        delete[] build_flags;
    }
    build_flags = new uint32_t[this->ref_cnt];
    for (int i = 0; i < this->ref_cnt; i++)
    {
        // TODO: specify build flags
        build_flags[i] = OPTIX_GEOMETRY_FLAG_DISABLE_ANYHIT;
    }

    // set build input
    build_input.type = OPTIX_BUILD_INPUT_TYPE_TRIANGLES;
    build_input.triangleArray.flags = build_flags;

    // vertex
    build_input.triangleArray.numVertices = this->vertexes.size();
    build_input.triangleArray.vertexFormat = OPTIX_VERTEX_FORMAT_FLOAT3;
    build_input.triangleArray.vertexBuffers =
        reinterpret_cast<CUdeviceptr *>(&(this->d_vertexes.d_ptr));
    build_input.triangleArray.vertexStrideInBytes = sizeof(Vertex);

    // index
    build_input.triangleArray.numIndexTriplets = this->triangles.size();
    build_input.triangleArray.indexFormat = OPTIX_INDICES_FORMAT_UNSIGNED_INT3;
    build_input.triangleArray.indexBuffer = this->d_triangles.d_pointer();
    build_input.triangleArray.indexStrideInBytes = sizeof(Triangle);

    // sbt
    uint sbt_idx_offset = 0;
    this->d_sbt_idx_offsets.resize(sizeof(uint));
    this->d_sbt_idx_offsets.upload(&sbt_idx_offset, sizeof(uint));
    build_input.triangleArray.numSbtRecords = 1;
    build_input.triangleArray.sbtIndexOffsetBuffer = d_sbt_idx_offsets.d_pointer();
    build_input.triangleArray.sbtIndexOffsetStrideInBytes = sizeof(uint);
    build_input.triangleArray.sbtIndexOffsetSizeInBytes = sizeof(uint);

    return build_input;
}

OptixAccelBuildOptions TriangleMesh::getASBuildOptions()
{
    OptixAccelBuildOptions build_options = {};

    build_options.buildFlags = OPTIX_BUILD_FLAG_ALLOW_COMPACTION;
    build_options.operation = OPTIX_BUILD_OPERATION_BUILD;

    return build_options;
}

void TriangleMesh::uploadVirtualProgramPayload(
    std::vector<std::shared_ptr<VirtualProgramEntry>> instances)
{
    const size_t n_instances = instances.size();
    const size_t data_size = sizeof(TriangleMeshData);
    const size_t memory_size = n_instances * data_size;

    // collect data
    TriangleMeshData *tmpdata = new TriangleMeshData[n_instances];
    for (size_t i = 0; i < n_instances; i++)
    {
        std::shared_ptr<TriangleMesh> m = std::dynamic_pointer_cast<TriangleMesh>(instances[i]);
        tmpdata[i].num_total_vertex = m->vertexes.size();
        tmpdata[i].vertexes = reinterpret_cast<Vertex *>(m->d_vertexes.d_pointer());
        tmpdata[i].num_total_triangle = m->triangles.size();
        tmpdata[i].triangles = reinterpret_cast<Triangle *>(m->d_triangles.d_pointer());
    }

    // upload data and free tmp
    TriangleMesh::sbt_data_payload.resize(memory_size);
    TriangleMesh::sbt_data_payload.upload(tmpdata, memory_size);
    delete[] tmpdata;
}

CUdeviceptr TriangleMesh::getPayloadDevicePtr()
{
    return TriangleMesh::sbt_data_payload.d_pointer();
}

void TriangleMesh::setProgram(std::shared_ptr<CudaProgram> prog)
{
    TriangleMesh::program = prog;
}

std::shared_ptr<CudaProgram> TriangleMesh::getProgram()
{
    return TriangleMesh::program;
}

OptixProgramGroupOptions &TriangleMesh::getProgramGroupOptions()
{
    return this->program_opt;
}

OptixProgramGroupDesc &TriangleMesh::getProgramGroupDesc()
{
    this->program_desc.hitgroup.moduleCH = this->cuda_module->optix_module;
    return this->program_desc;
}

void TriangleMesh::_onBuildEnd()
{
    this->updateAABB();
}

void TriangleMesh::updateAABB()
{
    // calc AABB
    this->aabb = Bounds::fromPoint(this->vertexes[0].pos);
    for (size_t i = 1; i < this->vertexes.size(); i++)
    {
        OptixAabb p = Bounds::fromPoint(this->vertexes[i].pos);
        this->aabb = Bounds::combine(this->aabb, p);
    }
    // upload aabb
    this->d_aabb.resize(sizeof(OptixAabb));
    this->d_aabb.upload(&this->aabb, sizeof(OptixAabb));
}